Repository navigation
fix(worldmodel): persist the EMA target encoder; pin the random baseline - #54
Merged
Merged
Conversation
First two items of the benchmark repair, plus a correction to node 14 that the work surfaced. 1. CHECKPOINT FORMAT v2 — persist target_encoder_state_dict. v1 saved encoder_state_dict and predictor_state_dict only. The failure mode was documented here, in frame_diagnostic.py and in ledger node 14 as "a restored model has target == encoder, so it reports zero mismatch". That is WRONG, and the truth is worse: __init__ deep-copies the target from the encoder AT CONSTRUCTION, before any load, so a v1 restore leaves the target at its RANDOM INITIALIZATION. Measured after loading encoder_state_dict into a fresh module: vs that module's own random init 0.00e+00 <- byte-identical vs the loaded online encoder 2.04 vs the trained target 0.39 target == encoder holds only if you construct and never load, which is not what using a checkpoint means. A v1-restored model plans toward goal latents from an UNTRAINED network -- not the one-frame remedy arm, and not any condition that was ever run. v2 restore verified exact (max parameter diff 0.00e+00). This matters beyond tidiness: the registered acceptance criterion for the benchmark repair says it must separate zero from random from MSE ON EXISTING DATA. Run against v1 checkpoints that would have been scoring an untrained goal encoder, and it would have looked like it worked. train.py has no equivalent defect -- it saves model.state_dict(), which already covers every submodule. export_coreml.py's GoalEncoder export is therefore also unaffected; checked rather than assumed. 2. PIN THE RANDOM BASELINE — baseline_reps, default 100. The baseline touches no encoder, no predictor and no CEM, only _step with uniform actions, so precision costs env steps alone. It was estimated at one draw per episode -- a scalar with binomial SE ~0.10 at n=20 -- and re-drawn per seed, injecting its own sampling noise into every arm comparison as if it were model variance. Decomposing node 14's "0.15 to 0.45" swing, it is NOT episode-limited binomial noise alone. Same episode set, varying only the draw: sd 0.0750 at reps=1, 0.0054 at reps=100, 0.0022 at reps=500. Different episode sets with reps pinned at 100: sd 0.0853. So it is roughly half sampling noise and half genuine episode-set difficulty. Reps removes the first half; decoupling the episode seed from the model seed (next item) is required for the second. Kept episode-paired rather than moved to an independent episode set: a paired comparison needs P(random succeeds | THIS episode). Per-episode rates are returned so a later paired analysis can use them as a difficulty covariate. random_baseline_stderr is returned so a reader can see the pin held. THE MODEL ARM IS BIT-IDENTICAL. Verified against node 14's recorded ld=32 success_rate at seeds 0-4 (0.30, 0.55, 0.20, 0.35, 0.35): exact match on all five. starts/goal_pos are drawn before the loop and the CEM generator is separate, so moving the baseline estimation after the model loop cannot perturb it -- but that is an argument, and the check is the evidence. 3. Ledger gains a dated "Corrections to node 14" block rather than a retro-edit of node 14 itself, matching the file's existing "Review clarifications" precedent. Node 14's measurements stand: forward_eval.py has no torch.load and trains in-process, so nothing it reported came from a restored checkpoint. Verified: smoke test passes on three consecutive runs; ledger gate, ADR gate and both test suites green. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DJJ8y1zo7WSqEmGNW8DQap
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
First two items of the benchmark repair — plus a correction to node 14 that doing the work surfaced.
1. Checkpoint format v2 — persist
target_encoder_state_dictThe documented failure mode was wrong, and the truth is worse. Node 14,
frame_diagnostic.pyand this file all said a restored model hastarget == encoder, so it reports zero mismatch. But__init__deep-copies the target from the encoder at construction, before any load — so a v1 restore leaves the target at its random initialization.Measured, after loading
encoder_state_dictinto a fresh module:target == encoderholds only if you construct and never load, which is not what using a checkpoint means. A v1-restored model plans toward goal latents from an untrained network — not the one-frame remedy arm, and not any condition that was ever run. v2 restore verified exact (0.00e+00 against the original target).This is not tidiness. The registered acceptance criterion says the repair must separate zero from random from MSE on existing data. Run against v1 checkpoints, it would have been scoring an untrained goal encoder — and it would have looked like it worked.
train.pyhas no equivalent defect: it savesmodel.state_dict(), which already covers every submodule. Soexport_coreml.py'sGoalEncoder.mlpackageis unaffected — checked, not assumed.2. Pin the random baseline —
baseline_reps, default 100The baseline touches no encoder, no predictor and no CEM, only
_stepwith uniform actions, so precision costs env steps alone. It was estimated at one draw per episode and re-drawn per seed, injecting its own sampling noise into every arm comparison as if it were model variance.Decomposing node 14's
0.15 → 0.45swing — it is not episode-limited binomial noise alone:Roughly half sampling noise, half genuine episode-set difficulty, since
starts/goal_posare redrawn fromrandom.Random(seed). Reps removes the first half; decoupling the episode seed from the model seed is required for the second.Kept episode-paired rather than moved to an independent set — a paired comparison needs
P(random succeeds | THIS episode). Per-episode rates are returned so a later paired analysis can use them as a difficulty covariate, andrandom_baseline_stderris returned so a reader can see the pin held.The model arm is bit-identical
Verified against node 14's recorded
ld=32success_rateat seeds 0–4 —0.30, 0.55, 0.20, 0.35, 0.35— exact match on all five.starts/goal_posare drawn before the loop and the CEM generator is separate, so moving baseline estimation after the model loop cannot perturb it. But that is an argument; the check is the evidence.3. Ledger correction, appended not retro-edited
A dated
## Corrections to node 14block, matching the file's existing## Review clarificationsprecedent. Node 14's measurements stand —forward_eval.pyhas notorch.loadand trains in-process, so nothing it reported came from a restored checkpoint. Only the parenthetical account of the hazard was wrong.Verification
Remaining in the ordering
Decouple episode seed from model seed → per-episode CEM generator → paired continuous primary → raise n to ~200. Not in this PR.